#include "ar_opencv.h"

const double akaze_thresh = 3e-4; // AKAZE detection threshold set to locate about 1000 keypoints
const double ransac_thresh = 2.5f; // RANSAC inlier threshold
const double nn_match_ratio = 0.8f; // Nearest-neighbour matching ratio
const int bb_min_inliers = 20; // Minimal number of inliers to draw bounding box

static Ptr<AKAZE> akaze;
static vector<KeyPoint> object_kp;
static Mat object_desc;
static vector<Point2f> object_bb(4);
// 暴力匹配
static Ptr<DescriptorMatcher> matcher = DescriptorMatcher::create("BruteForce-Hamming");

vector<Point2f> Points(vector<KeyPoint> keypoints) {
    vector<Point2f> res;
    for (unsigned i = 0; i < keypoints.size(); i++) {
        res.push_back(keypoints[i].pt);
    }
    return res;
}

long getCurrentTime() {
    struct timeval tv;
    gettimeofday(&tv,NULL);
    return tv.tv_sec * 1000 + tv.tv_usec / 1000;
}

// 匹配
vector<Point2f> process(const vector<KeyPoint>& kp, const Mat& desc) {
    vector< vector<DMatch> > matches;
    vector<KeyPoint> matched1, matched2;
    matcher->knnMatch(object_desc, desc, matches, 2);

    for(unsigned i = 0; i < matches.size(); i++) {
        if(matches[i][0].distance < nn_match_ratio * matches[i][1].distance) {
            matched1.push_back(object_kp[matches[i][0].queryIdx]);
            matched2.push_back(       kp[matches[i][0].trainIdx]);
        }
    }

    Mat inlier_mask, homography;
    vector<KeyPoint> inliers1, inliers2;
    vector<DMatch> inlier_matches;
    if (matched1.size() >= 4) {
        homography = findHomography(Points(matched1), Points(matched2), RANSAC, ransac_thresh, inlier_mask);
    }

    vector<Point2f> new_bb;
    if (matched1.size() >= 4 && !homography.empty()) {
        for (unsigned i = 0; i < matched1.size(); i++) {
            if (inlier_mask.at<uchar>(i)) {
                int new_i = static_cast<int>(inliers1.size());
                inliers1.push_back(matched1[i]);
                inliers2.push_back(matched2[i]);
                inlier_matches.push_back(DMatch(new_i, new_i, 0));
            }
        }

        //ratio = stats.inliers * 1.0 / stats.matches;

        perspectiveTransform(object_bb, new_bb, homography);

        if (inliers1.size() >= bb_min_inliers) {
            LOGE("inliers enough %d -----------", inliers1.size());
        } else {
            new_bb.clear();
        }
    }
    return new_bb;
}

// 初始化AR
void arInit(const Mat& frame) {
    akaze = AKAZE::create();
    akaze->setThreshold(akaze_thresh);

    akaze->detectAndCompute(frame, noArray(), object_kp, object_desc);

    object_bb[0] = Point(0, 0);
    object_bb[1] = Point(frame.cols, 0);
    object_bb[2] = Point(frame.cols, frame.rows);
    object_bb[3] = Point(0, frame.rows);

    LOGE("arInit : object kp %d", object_kp.size());
}

// AR追踪
bool arTracking(const Mat& frame, int width, int height, int (&bounds)[8]) {

    if (akaze == NULL) {
        return false;
    }

    LOGD("arTracking width %d, height %d", width, height);
    long first = getCurrentTime();

    vector<KeyPoint> kp;
    Mat desc;
    akaze->detectAndCompute(frame, noArray(), kp, desc);

    long second = getCurrentTime();
    LOGE("arTracking : detect time %ld, frame kp %d", second - first, kp.size());

    // 处理黑屏等特殊画面
    if (kp.size() < 100) {
        return false;
    }

    vector<Point2f> bb = process(kp, desc);

    long third = getCurrentTime();
    LOGE("arTracking : process time %ld", third - second);

    if (bb.size() != 4) {
        return false;
    }

    // 输出边缘点
    for (int i= 0 ; i < 4; ++i) {
        bounds[2*i    ] = bb[i].y;
        bounds[2*i + 1] = bb[i].x;
    }

    return true;
}
